import pandas as pd
import numpy as np
from castle.metrics import MetricsDAG
from castle.algorithms import TTPM

"""
dataset:是官网下载的数据集，与代码的datasets_phase2有对应关系，24对应1，25对应2，18无对应关系且节点较少为55
"""

if __name__ == '__main__':
    file_index = 18
    path = './dataset/'
    alarm_data = pd.read_csv(path + str(file_index) + '/Alarm.csv', encoding='utf')
    alarm_data = alarm_data.iloc[:, 0:3]
    alarm_data.columns = ['event', 'node', 'timestamp']
    alarm_data = alarm_data.reindex(columns=['event', 'timestamp', 'node'])
    all_event = np.sort(alarm_data['event'].unique())
    topology_matrix = np.zeros(shape=(alarm_data["node"].max() + 1, alarm_data["node"].max() + 1))
    est_causal_matrix = np.zeros(shape=(alarm_data["event"].max() + 1, alarm_data["event"].max() + 1), dtype=int)
    ttpm = TTPM(topology_matrix, max_hop=2, max_iter=100, delta=0.01, penalty='BIC', epsilon=0.6)

    for k in range(13):
        slice_data = alarm_data.iloc[k * 2500:k * 2500 + 5000, :].copy()
        ttpm.learn(slice_data)
        local_causal_matrix = np.array(ttpm.causal_matrix).astype(int)
        slice_event = np.sort(slice_data['event'].unique())
        missing_event = np.setdiff1d(all_event, slice_event)
        for i in missing_event:
            # 插入全零的新行
            new_row = np.zeros((1, local_causal_matrix.shape[1]), dtype=int)
            local_causal_matrix = np.concatenate([local_causal_matrix[:i], new_row, local_causal_matrix[i:]])
            # 插入全零的新列
            new_col = np.zeros((local_causal_matrix.shape[0], 1), dtype=int)
            local_causal_matrix = np.concatenate([local_causal_matrix[:, :i], new_col, local_causal_matrix[:, i:]],
                                                 axis=1)

        est_causal_matrix = est_causal_matrix | local_causal_matrix

    np.fill_diagonal(est_causal_matrix, 0)
    np.save(path + str(file_index) + '/est.npy', est_causal_matrix)
